import os
from PIL import Image
import numpy as np

from torch.utils import data

from data.augmentation import CDDataAugmentation

import pandas as pd


class GaofenDataset(data.Dataset):
    def __init__(
        self, root_dir, split="train", img_size=512, to_tensor=True, in_channels=3
    ):
        super(GaofenDataset, self).__init__()
        self.root_dir = root_dir
        self.to_tensor = to_tensor
        self.in_channels = in_channels
        csv_file = os.path.join(root_dir, f"cd_{split}.csv")
        self.csv_file = pd.read_csv(csv_file, header=None)
        if split == "train" or split == "train_all":
            self.augm = CDDataAugmentation(
                img_size=img_size,
                in_channels=in_channels,
                with_random_hflip=True,
                with_random_vflip=True,
                with_random_rot=True,
                with_scale_random_crop=True,
                with_random_blur=True,
                with_color=False,
            )
        else:
            self.augm = CDDataAugmentation(img_size=img_size, in_channels=in_channels)

    def __getitem__(self, index):
        filename = self.csv_file.iloc[index, 0]
        img1_path = os.path.join(self.root_dir, "image", f"{filename}_1.png")
        img2_path = os.path.join(self.root_dir, "image", f"{filename}_2.png")
        label_change_path = os.path.join(self.root_dir, "gt", f"{filename}_change.png")
        label_img1_path = os.path.join(self.root_dir, "gt", f"{filename}_1_label.png")
        label_img2_path = os.path.join(self.root_dir, "gt", f"{filename}_2_label.png")

        img1 = np.asarray(Image.open(img1_path).convert("RGB"))
        img2 = np.asarray(Image.open(img2_path).convert("RGB"))
        label_change = np.array(Image.open(label_change_path), dtype=np.uint8)
        label_change = label_change // 255
        labels = [label_change]

        if self.in_channels == 4:
            label_img1 = np.asarray(Image.open(label_img1_path), dtype=np.uint8)
            label_img2 = np.asarray(Image.open(label_img2_path), dtype=np.uint8)
            label_img1 = label_img1 // 255
            label_img2 = label_img2 // 255
            labels = [label_change, label_img1, label_img2]

        [img, img_B], labels = self.augm.transform(
            [img1, img2], labels, to_tensor=self.to_tensor,
        )
        return {"name": filename, "A": img, "B": img_B, "L": labels[0].long()[0]}

    def __len__(self):
        return len(self.csv_file)
